#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
系统管理脚本
整合了数据库初始化、用户权限设置、系统检查等功能
"""

import os
import sys
import time
import signal
import argparse
import subprocess
import requests
import re
import psutil
from datetime import datetime
from flask import Flask
from werkzeug.security import generate_password_hash
# 设置默认密码哈希方法为pbkdf2:sha256，避免使用scrypt（某些系统不支持）
from functools import partial
generate_password_hash = partial(generate_password_hash, method='pbkdf2:sha256')

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 导入模型
try:
    from models import db, User, Script, CustomScript, ScheduledTask, CustomTask
except ImportError:
    print("❌ 无法导入模型，请确保models.py文件存在")
    sys.exit(1)

# 创建Flask应用
def create_app():
    """创建Flask应用"""
    app = Flask(__name__)
    app.config['SECRET_KEY'] = 'system_management_secret_key'
    
    # 使用外部MySQL数据库
    from config.mysql_config import MySQLConfig
    mysql_config = MySQLConfig()
    mysql_config.validate_config()
    
    app.config['SQLALCHEMY_DATABASE_URI'] = mysql_config.get_database_uri()
    app.config['SQLALCHEMY_ENGINE_OPTIONS'] = mysql_config.get_engine_options()
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
    
    db.init_app(app)
    return app

# 数据库初始化功能
def init_database(args):
    """初始化数据库"""
    app = create_app()
    
    with app.app_context():
        print("开始初始化数据库...")
        
        try:
            # 创建所有表
            db.create_all()
            print("✅ 数据库表创建完成")
            
            # 检查是否已存在管理员用户
            admin_user = User.query.filter_by(username='admin').first()
            if admin_user:
                print("⚠️  管理员用户已存在，跳过创建")
                if args.force:
                    print("强制更新管理员密码...")
                    admin_user.password = generate_password_hash('admin')
                    db.session.commit()
                    print("✅ 管理员密码已重置为: admin")
            else:
                # 创建管理员用户
                admin = User(
                    username='admin',
                    password=generate_password_hash('admin'),
                    role='admin',
                    can_use_dingtalk=True,
                    can_use_custom_tasks=True,
                    can_manage_packages=True,
                    can_manage_users=True
                )
                db.session.add(admin)
                db.session.commit()
                print("✅ 管理员用户创建成功")
                print("用户名: admin")
                print("密码: admin")
                print("⚠️  请登录后立即修改默认密码!")
            
            return True
        except Exception as e:
            print(f"❌ 数据库初始化失败: {str(e)}")
            return False

# 用户权限初始化功能
def init_user_permissions(args):
    """初始化用户权限"""
    app = create_app()
    
    with app.app_context():
        # 获取所有用户
        users = User.query.all()
        
        print(f"找到 {len(users)} 个用户")
        
        for user in users:
            print(f"处理用户: {user.username}")
            
            # 设置默认权限
            if user.username == 'admin':
                # admin用户拥有所有权限
                user.role = 'admin'
                user.is_admin = True
                user.can_use_dingtalk = True
                user.can_use_custom_tasks = True
                user.can_manage_packages = True
                user.can_manage_users = True
                print(f"  - {user.username}: 设置为管理员，拥有所有权限")
            else:
                # 普通用户默认权限
                if user.role is None or user.role == '':
                    user.role = 'user'
                
                # 如果权限字段为None，设置默认值
                if user.can_use_dingtalk is None:
                    user.can_use_dingtalk = True
                if user.can_use_custom_tasks is None:
                    user.can_use_custom_tasks = args.custom_tasks
                if user.can_manage_packages is None:
                    user.can_manage_packages = False
                if user.can_manage_users is None:
                    user.can_manage_users = False
                
                print(f"  - {user.username}: 设置为普通用户，基本权限")
        
        # 保存更改
        db.session.commit()
        print("✅ 用户权限初始化完成")
        return True

# 检查数据库功能
def check_database(args):
    """检查数据库表结构"""
    app = create_app()
    
    with app.app_context():
        try:
            from sqlalchemy import text, inspect
            
            # 测试数据库连接
            db.session.execute(text('SELECT 1'))
            print("✅ MySQL数据库连接成功")
            
            # 获取数据库检查器
            inspector = inspect(db.engine)
            
            # 获取所有表
            tables = inspector.get_table_names()
            
            print("=== 数据库表列表 ===")
            for table in tables:
                print(f"- {table}")
            
            # 检查特定表
            if args.table:
                table_name = args.table
                print(f"\n=== 检查表: {table_name} ===")
                
                # 检查表是否存在
                if table_name not in tables:
                    print(f"❌ 表 {table_name} 不存在")
                    return False
                
                # 获取表结构
                columns = inspector.get_columns(table_name)
                print("表结构:")
                for col in columns:
                    print(f"  - {col['name']} ({col['type']})")
                
                # 获取表中的数据数量
                result = db.session.execute(text(f"SELECT COUNT(*) FROM {table_name}"))
                count = result.scalar()
                print(f"数据行数: {count}")
                
                # 如果有数据，显示前几行
                if count > 0:
                    result = db.session.execute(text(f"SELECT * FROM {table_name} LIMIT 5"))
                    rows = result.fetchall()
                    print("前5行数据:")
                    for row in rows:
                        print(f"  {row}")
            
            return True
            
        except Exception as e:
            print(f"❌ 数据库检查失败: {str(e)}")
            return False

# 检查自定义脚本功能
def check_custom_scripts(args):
    """检查自定义脚本"""
    app = create_app()
    
    with app.app_context():
        print("=== 检查自定义脚本 ===")
        
        # 获取所有自定义脚本
        scripts = CustomScript.query.all()
        print(f"总共有 {len(scripts)} 个自定义脚本:")
        
        for script in scripts:
            print(f"ID: {script.id}")
            print(f"名称: {script.name}")
            print(f"文件路径: {script.file_path}")
            print(f"用户ID: {script.user_id}")
            print(f"创建时间: {script.created_at}")
            print(f"文件是否存在: {os.path.exists(script.file_path) if script.file_path else False}")
            print("-" * 50)
        
        # 检查特定ID的脚本
        if args.id:
            script = CustomScript.query.get(args.id)
            if script:
                print(f"\nID为{args.id}的脚本存在:")
                print(f"名称: {script.name}")
                print(f"文件路径: {script.file_path}")
                print(f"用户ID: {script.user_id}")
            else:
                print(f"\nID为{args.id}的脚本不存在")
        
        return True

# 检查登录功能
def check_login(args):
    """检查登录功能"""
    # 创建会话
    session = requests.Session()
    
    # 获取登录页面
    try:
        login_page = session.get(f'http://localhost:{args.port}/login')
        print(f'登录页面状态码: {login_page.status_code}')
        
        if login_page.status_code != 200:
            print(f"❌ 无法访问登录页面，请确保Web服务已启动并监听在端口 {args.port}")
            return False
        
        # 打印登录页面的表单部分
        form_match = re.search(r'<form[^>]*action="[^"]*"[^>]*>(.*?)</form>', login_page.text, re.DOTALL)
        if form_match:
            print(f'找到登录表单')
        else:
            print('❌ 未找到登录表单')
            return False
        
        # 尝试登录 - 使用正确的表单字段
        login_data = {'username': args.username, 'password': args.password}
        
        # 检查是否有隐藏字段
        hidden_fields = re.findall(r'<input type="hidden" name="([^"]+)" value="([^"]+)"', login_page.text)
        for name, value in hidden_fields:
            print(f'隐藏字段: {name} = {value}')
            login_data[name] = value
        
        response = session.post(f'http://localhost:{args.port}/login', data=login_data, allow_redirects=True)
        print(f'登录状态码: {response.status_code}')
        print(f'登录后URL: {response.url}')
        
        # 检查是否登录成功 - 查找登出链接
        if 'logout' in response.text.lower():
            print('✅ 登录成功!')
            
            # 尝试访问仪表板
            dashboard_response = session.get(f'http://localhost:{args.port}/extended_dashboard')
            print(f'仪表板状态码: {dashboard_response.status_code}')
            
            # 检查是否成功访问仪表板
            if 'extended_dashboard' in dashboard_response.url:
                print('✅ 成功访问仪表板!')
                return True
            else:
                print('❌ 访问仪表板失败!')
                return False
        else:
            print('❌ 登录失败!')
            return False
    
    except requests.exceptions.ConnectionError:
        print(f"❌ 连接失败，请确保Web服务已启动并监听在端口 {args.port}")
        return False

# 检查调度器功能
def check_scheduler(args):
    """检查调度器状态"""
    app = create_app()
    
    with app.app_context():
        try:
            # 尝试导入调度器
            sys.path.append(os.path.join(os.path.dirname(__file__), 'services'))
            from scheduler_service import get_scheduler_status
            
            status = get_scheduler_status()
            if status:
                print("✅ 调度器状态检查成功")
                print(status)
                return True
            else:
                print("❌ 无法获取调度器状态")
                return False
        except ImportError:
            print("❌ 无法导入调度器服务模块，请确保scheduler_service.py文件存在")
            return False
        except Exception as e:
            print(f"❌ 检查调度器状态失败: {str(e)}")
            return False

# 重启执行器功能
def restart_executor(args):
    """重启执行器服务"""
    print("🔄 正在重启执行器服务...")
    
    # 查找执行器服务进程
    executor_process = None
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            # 检查命令行参数中是否包含指定名称
            if proc.info['cmdline'] and any('executor_service.py' in cmd for cmd in proc.info['cmdline']):
                executor_process = proc
                break
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    
    if executor_process:
        print(f"✅ 找到执行器服务进程 (PID: {executor_process.pid})")
        
        # 终止进程
        print("🛑 正在停止执行器服务...")
        try:
            executor_process.terminate()
            executor_process.wait(timeout=10)
            print("✅ 执行器服务已停止")
        except Exception as e:
            print(f"⚠️ 停止执行器服务失败: {str(e)}")
            try:
                executor_process.kill()
                print("⚠️ 执行器服务被强制停止")
            except Exception as e:
                print(f"❌ 强制停止执行器服务失败: {str(e)}")
                return False
    else:
        print("⚠️ 未找到执行器服务进程")
    
    # 启动执行器服务
    print("🚀 正在启动执行器服务...")
    try:
        executor_path = os.path.join(os.path.dirname(__file__), 'services', 'executor_service.py')
        if not os.path.exists(executor_path):
            print(f"❌ 执行器服务脚本不存在: {executor_path}")
            return False
        
        # 使用subprocess启动执行器服务
        subprocess.Popen(
            [sys.executable, executor_path, '--host', '0.0.0.0'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            start_new_session=True
        )
        
        print("✅ 执行器服务启动命令已发送")
        time.sleep(2)  # 等待服务启动
        
        # 检查服务是否成功启动
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if proc.info['cmdline'] and any('executor_service.py' in cmd for cmd in proc.info['cmdline']):
                    print(f"✅ 执行器服务已成功启动 (PID: {proc.pid})")
                    return True
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
        
        print("❌ 执行器服务启动失败")
        return False
    except Exception as e:
        print(f"❌ 启动执行器服务失败: {str(e)}")
        return False

# 应用数据库迁移功能
def apply_migration(args):
    """应用数据库迁移 - MySQL版本"""
    app = create_app()
    
    with app.app_context():
        print("开始应用数据库迁移...")
        
        try:
            from sqlalchemy import text, inspect
            
            # 测试数据库连接
            db.session.execute(text('SELECT 1'))
            print("✅ MySQL数据库连接成功")
            
            # 获取数据库检查器
            inspector = inspect(db.engine)
            
            # 根据迁移类型执行不同的迁移
            if args.type == 'custom_task_constraint':
                # 检查custom_task表是否存在
                if 'custom_task' not in inspector.get_table_names():
                    print("❌ custom_task表不存在，请先运行初始化数据库命令")
                    return False
                
                # 检查是否已经存在唯一约束
                indexes = inspector.get_indexes('custom_task')
                unique_constraint_exists = any(
                    idx.get('unique', False) and 
                    set(idx.get('column_names', [])) == {'name', 'user_id'}
                    for idx in indexes
                )
                
                if unique_constraint_exists:
                    print("✅ CustomTask表已存在唯一约束，无需迁移")
                    return True
                
                # 添加唯一约束
                print("添加唯一约束...")
                db.session.execute(text("""
                    ALTER TABLE custom_task 
                    ADD CONSTRAINT unique_name_user UNIQUE (name, user_id)
                """))
                
                db.session.commit()
                print("✅ CustomTask表唯一约束迁移完成")
            
            elif args.type == 'retry_fields':
                # 检查scheduled_task表是否存在
                if 'scheduled_task' not in inspector.get_table_names():
                    print("❌ scheduled_task表不存在，请先运行初始化数据库命令")
                    return False
                
                # 检查scheduled_task表是否已经存在重试相关字段
                columns = inspector.get_columns('scheduled_task')
                column_names = [col['name'] for col in columns]
                
                if 'max_retries' in column_names and 'retry_delay' in column_names:
                    print("✅ ScheduledTask表已存在重试相关字段，无需迁移")
                    return True
                
                # 添加重试相关字段
                if 'max_retries' not in column_names:
                    print("添加max_retries字段...")
                    db.session.execute(text("""
                        ALTER TABLE scheduled_task 
                        ADD COLUMN max_retries INTEGER DEFAULT 0
                    """))
                
                if 'retry_delay' not in column_names:
                    print("添加retry_delay字段...")
                    db.session.execute(text("""
                        ALTER TABLE scheduled_task 
                        ADD COLUMN retry_delay INTEGER DEFAULT 60
                    """))
                
                if 'retry_count' not in column_names:
                    print("添加retry_count字段...")
                    db.session.execute(text("""
                        ALTER TABLE scheduled_task 
                        ADD COLUMN retry_count INTEGER DEFAULT 0
                    """))
                
                if 'last_error' not in column_names:
                    print("添加last_error字段...")
                    db.session.execute(text("""
                        ALTER TABLE scheduled_task 
                        ADD COLUMN last_error TEXT
                    """))
                
                db.session.commit()
                print("✅ ScheduledTask表重试相关字段迁移完成")
            
            else:
                print(f"❌ 未知的迁移类型: {args.type}")
                return False
            
            return True
        
        except Exception as e:
            print(f"❌ 应用迁移失败: {str(e)}")
            db.session.rollback()
            return False

# 主函数
def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='系统管理工具')
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # 初始化数据库命令
    init_db_parser = subparsers.add_parser('init_db', help='初始化数据库')
    init_db_parser.add_argument('--force', action='store_true', help='强制重置管理员密码')
    
    # 初始化用户权限命令
    init_permissions_parser = subparsers.add_parser('init_permissions', help='初始化用户权限')
    init_permissions_parser.add_argument('--custom-tasks', action='store_true', help='允许普通用户使用自定义任务')
    
    # 检查数据库命令
    check_db_parser = subparsers.add_parser('check_db', help='检查数据库')
    check_db_parser.add_argument('--table', help='指定要检查的表名')
    
    # 检查自定义脚本命令
    check_scripts_parser = subparsers.add_parser('check_scripts', help='检查自定义脚本')
    check_scripts_parser.add_argument('--id', type=int, help='指定要检查的脚本ID')
    
    # 检查登录命令
    check_login_parser = subparsers.add_parser('check_login', help='检查登录功能')
    check_login_parser.add_argument('--port', type=int, default=8088, help='Web服务端口')
    check_login_parser.add_argument('--username', default='admin', help='登录用户名')
    check_login_parser.add_argument('--password', default='admin', help='登录密码')
    
    # 检查调度器命令
    check_scheduler_parser = subparsers.add_parser('check_scheduler', help='检查调度器状态')
    
    # 重启执行器命令
    restart_executor_parser = subparsers.add_parser('restart_executor', help='重启执行器服务')
    
    # 应用迁移命令
    apply_migration_parser = subparsers.add_parser('apply_migration', help='应用数据库迁移')
    apply_migration_parser.add_argument('--type', choices=['custom_task_constraint', 'retry_fields'], 
                                      required=True, help='迁移类型')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 执行对应的命令
    if args.command == 'init_db':
        init_database(args)
    elif args.command == 'init_permissions':
        init_user_permissions(args)
    elif args.command == 'check_db':
        check_database(args)
    elif args.command == 'check_scripts':
        check_custom_scripts(args)
    elif args.command == 'check_login':
        check_login(args)
    elif args.command == 'check_scheduler':
        check_scheduler(args)
    elif args.command == 'restart_executor':
        restart_executor(args)
    elif args.command == 'apply_migration':
        apply_migration(args)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()